# traffic_sign_model.py
import tensorflow as tf
import pandas as pd
import pathlib
import numpy as np

# 配置GPU显存按需分配
gpus = tf.config.list_physical_devices("GPU")
if gpus:
    tf.config.experimental.set_memory_growth(gpus[0], True)
    tf.config.set_visible_devices([gpus[0]], "GPU")

# 数据配置参数
DATA_DIR = "datasets/14_traffic_sign"
IMAGE_DIR = f"{DATA_DIR}/images"
ANNOTATION_FILE = f"{DATA_DIR}/annotations.csv"
BATCH_SIZE = 32
IMG_SIZE = (128, 128)  # 根据实际需求调整

# 数据预处理函数
def load_and_preprocess(path, label):
    image = tf.io.read_file(path)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, IMG_SIZE)
    return image/255.0, label

# 构建模型
def create_model(input_shape, num_classes):
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=input_shape),
        tf.keras.layers.MaxPooling2D((2,2)),
        tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
        tf.keras.layers.MaxPooling2D((2,2)),
        tf.keras.layers.Conv2D(128, (3,3), activation='relu'),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(256, activation='relu'),
        tf.keras.layers.Dense(num_classes)
    ])
    return model

def main():
    # 加载数据
    annotations = pd.read_csv(ANNOTATION_FILE)
    image_paths = [f"{IMAGE_DIR}/{name}" for name in annotations["file_name"].values]
    labels = annotations["category"].values

    # 创建数据集
    ds = tf.data.Dataset.from_tensor_slices((image_paths, labels))
    ds = ds.shuffle(len(image_paths))
    ds = ds.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
    ds = ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

    # 划分训练集和验证集
    train_size = int(0.8 * len(ds))
    train_ds = ds.take(train_size)
    val_ds = ds.skip(train_size)

    # 创建并编译模型
    model = create_model((*IMG_SIZE, 3), 58)  # 假设有43个类别
    model.compile(
        optimizer='adam',
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=['accuracy']
    )

    # 添加模型保存回调
    callbacks = [
        tf.keras.callbacks.ModelCheckpoint(
            'best_model.h5',
            save_best_only=True,
            monitor='val_accuracy'
        )
    ]

    # 训练模型
    history = model.fit(
        train_ds,
        validation_data=val_ds,
        epochs=20,
        callbacks=callbacks
    )

    # 保存最终模型
    model.save('traffic_sign_model.h5')
    print("模型已保存为 traffic_sign_model.h5")

if __name__ == "__main__":
    main()